package org.apache.lucene.search;

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import java.io.IOException;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.util.Counter;
import org.apache.lucene.util.ThreadInterruptedException;

/**
 * The {@link TimeLimitingCollector} is used to timeout search requests that
 * take longer than the maximum allowed search time limit. After this time is
 * exceeded, the search thread is stopped by throwing a
 * {@link TimeExceededException}.
 */
public class TimeLimitingCollector extends Collector {

    /** Thrown when elapsed search time exceeds allowed search time. */
    @SuppressWarnings("serial")
    public static class TimeExceededException extends RuntimeException {
        private long timeAllowed;
        private long timeElapsed;
        private int lastDocCollected;

        private TimeExceededException(long timeAllowed, long timeElapsed, int lastDocCollected) {
            super("Elapsed time: " + timeElapsed + "Exceeded allowed search time: " + timeAllowed + " ms.");
            this.timeAllowed = timeAllowed;
            this.timeElapsed = timeElapsed;
            this.lastDocCollected = lastDocCollected;
        }

        /** Returns allowed time (milliseconds). */
        public long getTimeAllowed() {
            return timeAllowed;
        }

        /** Returns elapsed time (milliseconds). */
        public long getTimeElapsed() {
            return timeElapsed;
        }

        /** Returns last doc (absolute doc id) that was collected when the search time exceeded. */
        public int getLastDocCollected() {
            return lastDocCollected;
        }
    }

    private long t0 = Long.MIN_VALUE;
    private long timeout = Long.MIN_VALUE;
    private final Collector collector;
    private final Counter clock;
    private final long ticksAllowed;
    private boolean greedy = false;
    private int docBase;

    /**
     * Create a TimeLimitedCollector wrapper over another {@link Collector} with a specified timeout.
     * @param collector the wrapped {@link Collector}
     * @param clock the timer clock
     * @param ticksAllowed max time allowed for collecting
     * hits after which {@link TimeExceededException} is thrown
     */
    public TimeLimitingCollector(final Collector collector, Counter clock, final long ticksAllowed) {
        this.collector = collector;
        this.clock = clock;
        this.ticksAllowed = ticksAllowed;
    }

    /**
     * Sets the baseline for this collector. By default the collectors baseline is 
     * initialized once the first reader is passed to the collector. 
     * To include operations executed in prior to the actual document collection
     * set the baseline through this method in your prelude.
     * <p>
     * Example usage:
     * <pre>
     *   Counter clock = ...;
     *   long baseline = clock.get();
     *   // ... prepare search
     *   TimeLimitingCollector collector = new TimeLimitingCollector(c, clock, numTicks);
     *   collector.setBaseline(baseline);
     *   indexSearcher.search(query, collector);
     * </pre>
     * </p>
     * @see #setBaseline() 
     * @param clockTime
     */
    public void setBaseline(long clockTime) {
        t0 = clockTime;
        timeout = t0 + ticksAllowed;
    }

    /**
     * Syntactic sugar for {@link #setBaseline(long)} using {@link Counter#get()}
     * on the clock passed to the construcutor.
     */
    public void setBaseline() {
        setBaseline(clock.get());
    }

    /**
     * Checks if this time limited collector is greedy in collecting the last hit.
     * A non greedy collector, upon a timeout, would throw a {@link TimeExceededException} 
     * without allowing the wrapped collector to collect current doc. A greedy one would 
     * first allow the wrapped hit collector to collect current doc and only then 
     * throw a {@link TimeExceededException}.
     * @see #setGreedy(boolean)
     */
    public boolean isGreedy() {
        return greedy;
    }

    /**
     * Sets whether this time limited collector is greedy.
     * @param greedy true to make this time limited greedy
     * @see #isGreedy()
     */
    public void setGreedy(boolean greedy) {
        this.greedy = greedy;
    }

    /**
     * Calls {@link Collector#collect(int)} on the decorated {@link Collector}
     * unless the allowed time has passed, in which case it throws an exception.
     * 
     * @throws TimeExceededException
     *           if the time allowed has exceeded.
     */
    @Override
    public void collect(final int doc) throws IOException {
        final long time = clock.get();
        if (timeout < time) {
            if (greedy) {
                //System.out.println(this+"  greedy: before failing, collecting doc: "+(docBase + doc)+"  "+(time-t0));
                collector.collect(doc);
            }
            //System.out.println(this+"  failing on:  "+(docBase + doc)+"  "+(time-t0));
            throw new TimeExceededException(timeout - t0, time - t0, docBase + doc);
        }
        //System.out.println(this+"  collecting: "+(docBase + doc)+"  "+(time-t0));
        collector.collect(doc);
    }

    @Override
    public void setNextReader(IndexReader reader, int base) throws IOException {
        collector.setNextReader(reader, base);
        this.docBase = base;
        if (Long.MIN_VALUE == t0) {
            setBaseline();
        }
    }

    @Override
    public void setScorer(Scorer scorer) throws IOException {
        collector.setScorer(scorer);
    }

    @Override
    public boolean acceptsDocsOutOfOrder() {
        return collector.acceptsDocsOutOfOrder();
    }

    /**
     * Returns the global TimerThreads {@link Counter}
     * <p>
     * Invoking this creates may create a new instance of {@link TimerThread} iff
     * the global {@link TimerThread} has never been accessed before. The thread
     * returned from this method is started on creation and will be alive unless
     * you stop the {@link TimerThread} via {@link TimerThread#stopTimer()}.
     * </p>
     * @return the global TimerThreads {@link Counter}
     * @lucene.experimental
     */
    public static Counter getGlobalCounter() {
        return TimerThreadHolder.THREAD.counter;
    }

    /**
     * Returns the global {@link TimerThread}.
     * <p>
     * Invoking this creates may create a new instance of {@link TimerThread} iff
     * the global {@link TimerThread} has never been accessed before. The thread
     * returned from this method is started on creation and will be alive unless
     * you stop the {@link TimerThread} via {@link TimerThread#stopTimer()}.
     * </p>
     * 
     * @return the global {@link TimerThread}
     * @lucene.experimental
     */
    public static TimerThread getGlobalTimerThread() {
        return TimerThreadHolder.THREAD;
    }

    private static final class TimerThreadHolder {
        static final TimerThread THREAD;
        static {
            THREAD = new TimerThread(Counter.newCounter(true));
            THREAD.start();
        }
    }

    /**
     * @lucene.experimental
     */
    public static final class TimerThread extends Thread {

        public static final String THREAD_NAME = "TimeLimitedCollector timer thread";
        public static final int DEFAULT_RESOLUTION = 20;
        // NOTE: we can avoid explicit synchronization here for several reasons:
        // * updates to volatile long variables are atomic
        // * only single thread modifies this value
        // * use of volatile keyword ensures that it does not reside in
        //   a register, but in main memory (so that changes are visible to
        //   other threads).
        // * visibility of changes does not need to be instantaneous, we can
        //   afford losing a tick or two.
        //
        // See section 17 of the Java Language Specification for details.
        private volatile long time = 0;
        private volatile boolean stop = false;
        private volatile long resolution;
        final Counter counter;

        public TimerThread(long resolution, Counter counter) {
            super(THREAD_NAME);
            this.resolution = resolution;
            this.counter = counter;
            this.setDaemon(true);
        }

        public TimerThread(Counter counter) {
            this(DEFAULT_RESOLUTION, counter);
        }

        @Override
        public void run() {
            while (!stop) {
                // TODO: Use System.nanoTime() when Lucene moves to Java SE 5.
                counter.addAndGet(resolution);
                try {
                    Thread.sleep(resolution);
                } catch (InterruptedException ie) {
                    throw new ThreadInterruptedException(ie);
                }
            }
        }

        /**
         * Get the timer value in milliseconds.
         */
        public long getMilliseconds() {
            return time;
        }

        /**
         * Stops the timer thread 
         */
        public void stopTimer() {
            stop = true;
        }

        /** 
         * Return the timer resolution.
         * @see #setResolution(long)
         */
        public long getResolution() {
            return resolution;
        }

        /**
         * Set the timer resolution.
         * The default timer resolution is 20 milliseconds. 
         * This means that a search required to take no longer than 
         * 800 milliseconds may be stopped after 780 to 820 milliseconds.
         * <br>Note that: 
         * <ul>
         * <li>Finer (smaller) resolution is more accurate but less efficient.</li>
         * <li>Setting resolution to less than 5 milliseconds will be silently modified to 5 milliseconds.</li>
         * <li>Setting resolution smaller than current resolution might take effect only after current 
         * resolution. (Assume current resolution of 20 milliseconds is modified to 5 milliseconds, 
         * then it can take up to 20 milliseconds for the change to have effect.</li>
         * </ul>      
         */
        public void setResolution(long resolution) {
            this.resolution = Math.max(resolution, 5); // 5 milliseconds is about the minimum reasonable time for a Object.wait(long) call.
        }
    }

}
